// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/parse/expressions.h>
#include <pollux/common/base/exceptions.h>
#include <pollux/core/expressions.h>
#include <pollux/exec/aggregate.h>
#include <pollux/expression/simple_function_registry.h>
#include <pollux/functions/function_registry.h>
#include <pollux/type/type.h>
#include <pollux/vector/constant_vector.h>
#include <pollux/vector/variant_to_vector.h>

namespace kumo::pollux::core {
    // static
    Expressions::TypeResolverHook Expressions::resolverHook_;
    // static
    Expressions::FieldAccessHook Expressions::fieldAccessHook_;

    namespace {
        // Determine output type based on input types.
        TypePtr resolveTypeImpl(
            std::vector<TypedExprPtr> inputs,
            const std::shared_ptr<const CallExpr> &expr,
            bool nullOnFailure) {
            POLLUX_CHECK_NOT_NULL(Expressions::getResolverHook());
            return Expressions::getResolverHook()(inputs, expr, nullOnFailure);
        }

        namespace {
            std::shared_ptr<const core::CastTypedExpr> makeTypedCast(
                const TypePtr &type,
                const TypedExprPtr &input) {
                return std::make_shared<const core::CastTypedExpr>(type, input, false);
            }

            std::vector<TypePtr> implicitCastTargets(const TypePtr &type) {
                std::vector<TypePtr> targetTypes;
                switch (type->kind()) {
                    // We decide not to implicitly upcast booleans because it maybe funky.
                    case TypeKind::BOOLEAN:
                        break;
                    case TypeKind::TINYINT:
                        targetTypes.emplace_back(SMALLINT());
                        [[fallthrough]];
                    case TypeKind::SMALLINT:
                        targetTypes.emplace_back(INTEGER());
                        targetTypes.emplace_back(REAL());
                        [[fallthrough]];
                    case TypeKind::INTEGER:
                        targetTypes.emplace_back(BIGINT());
                        targetTypes.emplace_back(DOUBLE());
                        [[fallthrough]];
                    case TypeKind::BIGINT:
                        break;
                    case TypeKind::REAL:
                        targetTypes.emplace_back(DOUBLE());
                        [[fallthrough]];
                    case TypeKind::DOUBLE:
                        break;
                    case TypeKind::ARRAY: {
                        auto childTargetTypes = implicitCastTargets(type->childAt(0));
                        for (auto childTarget: childTargetTypes) {
                            targetTypes.emplace_back(ARRAY(childTarget));
                        }
                        break;
                    }
                    default: // make compilers happy
                        break;
                }
                return targetTypes;
            }
        } // namespace

        // All acceptable implicit casts on this expression.
        // TODO: If we get this to be recursive somehow, we can save on cast function
        // signatures that need to be compiled and registered.
        std::vector<TypedExprPtr> genImplicitCasts(const TypedExprPtr &typedExpr) {
            auto targetTypes = implicitCastTargets(typedExpr->type());

            std::vector<TypedExprPtr> implicitCasts;
            implicitCasts.reserve(targetTypes.size());
            for (auto targetType: targetTypes) {
                implicitCasts.emplace_back(makeTypedCast(targetType, typedExpr));
            }
            return implicitCasts;
        }

        // TODO: Arguably all of this could be done with just Types.
        TypedExprPtr adjustLastNArguments(
            std::vector<TypedExprPtr> inputs,
            const std::shared_ptr<const CallExpr> &expr,
            size_t n) {
            auto type = resolveTypeImpl(inputs, expr, true /*nullOnFailure*/);
            if (type != nullptr) {
                return std::make_unique<CallTypedExpr>(
                    type, inputs, std::string{expr->getFunctionName()});
            }

            if (n == 0) {
                return nullptr;
            }

            size_t firstOfLastN = inputs.size() - n;
            // use it
            std::vector<TypedExprPtr> viableExprs{inputs[firstOfLastN]};
            // or lose it
            auto &&implicitCasts = genImplicitCasts(inputs[firstOfLastN]);
            std::move(
                implicitCasts.begin(),
                implicitCasts.end(),
                std::back_inserter(viableExprs));

            for (auto &viableExpr: viableExprs) {
                inputs[firstOfLastN] = viableExpr;
                auto adjustedExpr = adjustLastNArguments(inputs, expr, n - 1);
                if (adjustedExpr != nullptr) {
                    return adjustedExpr;
                }
            }

            return nullptr;
        }

        TypedExprPtr createWithImplicitCast(
            const std::shared_ptr<const core::CallExpr> &expr,
            const std::vector<TypedExprPtr> &inputs) {
            auto adjusted = adjustLastNArguments(inputs, expr, inputs.size());
            if (adjusted) {
                return adjusted;
            }
            auto type = resolveTypeImpl(inputs, expr, false /*nullOnFailure*/);
            return std::make_shared<CallTypedExpr>(
                type, std::move(inputs), std::string{expr->getFunctionName()});
        }

        bool isLambdaArgument(const exec::TypeSignature &typeSignature) {
            return typeSignature.baseName() == "function";
        }

        bool isLambdaArgument(const exec::TypeSignature &typeSignature, int numInputs) {
            return isLambdaArgument(typeSignature) &&
                   (typeSignature.parameters().size() == numInputs + 1);
        }

        bool hasLambdaArgument(const exec::FunctionSignature &signature) {
            for (const auto &type: signature.argumentTypes()) {
                if (isLambdaArgument(type)) {
                    return true;
                }
            }

            return false;
        }
    } // namespace

    // static
    TypedExprPtr Expressions::inferTypes(
        const std::shared_ptr<const core::IExpr> &expr,
        const TypePtr &inputRow,
        memory::MemoryPool *pool,
        const VectorPtr &complexConstants) {
        return inferTypes(expr, inputRow, {}, pool, complexConstants);
    }

    // static
    TypedExprPtr Expressions::inferTypes(
        const std::shared_ptr<const core::IExpr> &expr,
        const TypePtr &inputRow,
        const std::vector<TypePtr> &lambdaInputTypes,
        memory::MemoryPool *pool,
        const VectorPtr &complexConstants) {
        POLLUX_CHECK_NOT_NULL(expr);

        if (auto lambdaExpr = std::dynamic_pointer_cast<const LambdaExpr>(expr)) {
            return resolveLambdaExpr(
                lambdaExpr, inputRow, lambdaInputTypes, pool, complexConstants);
        }

        if (auto call = std::dynamic_pointer_cast<const CallExpr>(expr)) {
            if (!expr->getInputs().empty()) {
                if (auto returnType = tryResolveCallWithLambdas(
                    call, inputRow, pool, complexConstants)) {
                    return returnType;
                }
            }
        }

        // try rebuilding complex constant type from vector
        if (auto fun = std::dynamic_pointer_cast<const CallExpr>(expr)) {
            if (fun->getFunctionName() == "__complex_constant") {
                POLLUX_CHECK_NOT_NULL(
                    complexConstants,
                    "Expression contains __complex_constant function call, but complexConstants is missing");

                auto ccInputRow = complexConstants->as<RowVector>();
                POLLUX_CHECK_NOT_NULL(
                    ccInputRow,
                    "Expected RowVector for complexConstants: {}",
                    complexConstants->toString());
                auto name =
                        std::dynamic_pointer_cast<const FieldAccessExpr>(fun->getInputs()[0])
                        ->getFieldName();
                auto rowType = as_row_type(ccInputRow->type());
                return std::make_shared<const ConstantTypedExpr>(
                    ccInputRow->childAt(rowType->getChildIdx(name)));
            }
        }

        std::vector<TypedExprPtr> children;
        for (auto &child: expr->getInputs()) {
            children.push_back(
                inferTypes(child, inputRow, lambdaInputTypes, pool, complexConstants));
        }

        if (auto fae = std::dynamic_pointer_cast<const FieldAccessExpr>(expr)) {
            if (fieldAccessHook_) {
                auto result = fieldAccessHook_(fae, children);
                if (result) {
                    return result;
                }
            }
            POLLUX_CHECK(
                !fae->getFieldName().empty(), "Anonymous columns are not supported");
            POLLUX_CHECK_EQ(
                children.size(), 1, "Unexpected number of children in FieldAccessExpr");
            auto input = children.at(0)->type();
            auto &row = input->as_row();
            auto childIndex = row.getChildIdx(fae->getFieldName());
            if (fae->isRootColumn()) {
                return std::make_shared<FieldAccessTypedExpr>(
                    input->childAt(childIndex),
                    children.at(0),
                    std::string{fae->getFieldName()});
            } else {
                return std::make_shared<DereferenceTypedExpr>(
                    input->childAt(childIndex), children.at(0), childIndex);
            }
        }
        if (auto fun = std::dynamic_pointer_cast<const CallExpr>(expr)) {
            return createWithImplicitCast(fun, std::move(children));
        }
        if (auto input = std::dynamic_pointer_cast<const InputExpr>(expr)) {
            return std::make_shared<const InputTypedExpr>(inputRow);
        }
        if (auto constant = std::dynamic_pointer_cast<const ConstantExpr>(expr)) {
            if (constant->type()->kind() == TypeKind::ARRAY) {
                // Transform variant vector into an ArrayVector, then wrap it into a
                // ConstantVector<ComplexType>.
                POLLUX_CHECK_NOT_NULL(
                    pool, "parsing array literals requires a memory pool");
                VectorPtr constant_vector;
                if (constant->value().isNull()) {
                    constant_vector =
                            BaseVector::create_null_constant(constant->type(), 1, pool);
                } else {
                    auto array_vector = variantArrayToVector(
                        constant->type(), constant->value().array(), pool);
                    constant_vector = std::make_shared<ConstantVector<pollux::ComplexType> >(
                        pool, 1, 0, array_vector);
                }
                return std::make_shared<const ConstantTypedExpr>(constant_vector);
            }
            return std::make_shared<const ConstantTypedExpr>(
                constant->type(), constant->value());
        }
        if (auto cast = std::dynamic_pointer_cast<const CastExpr>(expr)) {
            return std::make_shared<const CastTypedExpr>(
                cast->type(), std::move(children), cast->nullOnFailure());
        }
        if (auto alreadyTyped = std::dynamic_pointer_cast<const ITypedExpr>(expr)) {
            return alreadyTyped;
        }

        POLLUX_FAIL("Unknown expression type: {}", expr->toString());
    }

    // static
    TypedExprPtr Expressions::resolveLambdaExpr(
        const std::shared_ptr<const core::LambdaExpr> &lambdaExpr,
        const TypePtr &inputRow,
        const std::vector<TypePtr> &lambdaInputTypes,
        memory::MemoryPool *pool,
        const VectorPtr &complexConstants) {
        auto names = lambdaExpr->inputNames();
        auto body = lambdaExpr->body();

        POLLUX_CHECK_LE(names.size(), lambdaInputTypes.size());
        std::vector<TypePtr> types;
        for (auto i = 0; i < names.size(); ++i) {
            types.push_back(lambdaInputTypes[i]);
        }

        auto signature =
                ROW(std::vector<std::string>(names), std::vector<TypePtr>(types));

        auto &inputRowType = inputRow->as_row();
        for (auto i = 0; i < inputRowType.size(); ++i) {
            if (!signature->containsChild(inputRowType.names()[i])) {
                names.push_back(inputRowType.names()[i]);
                types.push_back(inputRowType.childAt(i));
            }
        }

        auto lambdaRow = ROW(std::move(names), std::move(types));

        return std::make_shared<LambdaTypedExpr>(
            signature, inferTypes(body, lambdaRow, pool, complexConstants));
    }

    namespace {
        bool isLambdaSignature(
            const exec::FunctionSignature *signature,
            const std::shared_ptr<const CallExpr> &callExpr) {
            if (!hasLambdaArgument(*signature)) {
                return false;
            }

            const auto numArguments = callExpr->getInputs().size();

            if (numArguments != signature->argumentTypes().size()) {
                return false;
            }

            bool match = true;
            for (auto i = 0; i < numArguments; ++i) {
                if (auto lambda = dynamic_cast<const core::LambdaExpr *>(
                    callExpr->getInputs()[i].get())) {
                    const auto numLambdaInputs = lambda->inputNames().size();
                    const auto &argumentType = signature->argumentTypes()[i];
                    if (!isLambdaArgument(argumentType, numLambdaInputs)) {
                        match = false;
                        break;
                    }
                }
            }

            return match;
        }

        const exec::FunctionSignature *findLambdaSignature(
            const std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > &
            signatures,
            const std::shared_ptr<const CallExpr> &callExpr) {
            const exec::FunctionSignature *matchingSignature = nullptr;
            for (const auto &signature: signatures) {
                if (isLambdaSignature(signature.get(), callExpr)) {
                    POLLUX_CHECK_NULL(
                        matchingSignature,
                        "Cannot resolve ambiguous lambda function signatures for {}.",
                        callExpr->getFunctionName());
                    matchingSignature = signature.get();
                }
            }

            return matchingSignature;
        }

        const exec::FunctionSignature *findLambdaSignature(
            const std::vector<const exec::FunctionSignature *> &signatures,
            const std::shared_ptr<const CallExpr> &callExpr) {
            const exec::FunctionSignature *matchingSignature = nullptr;
            for (const auto &signature: signatures) {
                if (isLambdaSignature(signature, callExpr)) {
                    POLLUX_CHECK_NULL(
                        matchingSignature,
                        "Cannot resolve ambiguous lambda function signatures for {}.",
                        callExpr->getFunctionName());
                    matchingSignature = signature;
                }
            }

            return matchingSignature;
        }

        // Assumes no overlap in function names between scalar and aggregate functions,
        // i.e. 'foo' is either a scalar or aggregate function.
        const exec::FunctionSignature *findLambdaSignature(
            const std::shared_ptr<const CallExpr> &callExpr) {
            // Look for a scalar lambda function.
            auto scalarSignatures = get_function_signatures(callExpr->getFunctionName());
            if (!scalarSignatures.empty()) {
                return findLambdaSignature(scalarSignatures, callExpr);
            }

            // Look for an aggregate lambda function.
            if (auto signatures =
                    exec::getAggregateFunctionSignatures(callExpr->getFunctionName())) {
                return findLambdaSignature(signatures.value(), callExpr);
            }

            return nullptr;
        }
    } // namespace

    // static
    TypedExprPtr Expressions::tryResolveCallWithLambdas(
        const std::shared_ptr<const CallExpr> &callExpr,
        const TypePtr &inputRow,
        memory::MemoryPool *pool,
        const VectorPtr &complexConstants) {
        auto signature = findLambdaSignature(callExpr);

        if (signature == nullptr) {
            return nullptr;
        }

        // Resolve non-lambda arguments first.
        auto numArgs = callExpr->getInputs().size();
        std::vector<TypedExprPtr> children(numArgs);
        std::vector<TypePtr> childTypes(numArgs);
        for (auto i = 0; i < numArgs; ++i) {
            if (!isLambdaArgument(signature->argumentTypes()[i])) {
                children[i] = inferTypes(
                    callExpr->getInputs()[i], inputRow, pool, complexConstants);
                childTypes[i] = children[i]->type();
            }
        }

        // Resolve lambda arguments.
        exec::SignatureBinder binder(*signature, childTypes);
        binder.tryBind();
        for (auto i = 0; i < numArgs; ++i) {
            auto argSignature = signature->argumentTypes()[i];
            if (isLambdaArgument(argSignature)) {
                std::vector<TypePtr> lambdaTypes;
                for (auto j = 0; j < argSignature.parameters().size() - 1; ++j) {
                    auto type = binder.tryResolveType(argSignature.parameters()[j]);
                    if (type == nullptr) {
                        return nullptr;
                    }
                    lambdaTypes.push_back(type);
                }

                children[i] = inferTypes(
                    callExpr->getInputs()[i],
                    inputRow,
                    lambdaTypes,
                    pool,
                    complexConstants);
            }
        }

        return createWithImplicitCast(callExpr, std::move(children));
    }

    // This method returns null if the expression doesn't depend on any input row.
    TypePtr Expressions::getInputRowType(const TypedExprPtr &expr) {
        if (auto inputExpr = std::dynamic_pointer_cast<const InputTypedExpr>(expr)) {
            return inputExpr->type();
        }
        TypePtr inputRowType;
        for (auto &input: expr->inputs()) {
            auto childRowType = getInputRowType(input);
            if (childRowType) {
                POLLUX_USER_CHECK(!inputRowType || *inputRowType == *childRowType);
                inputRowType = childRowType;
            }
        }

        return inputRowType;
    }
} // namespace kumo::pollux::core
